package org.etlt.extract;

import org.etlt.EtltRuntimeException;
import org.etlt.job.JobContext;
import org.etlt.job.PkRange;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;

public class BigDatabaseExtractor extends DatabaseExtractor {

    private PreparedStatement pkStatement;

    private ResultSet pkResultSet;

    private transient Object minPk;

    private transient Object maxPk;

    private transient Object pkBoundary;

    private transient boolean complemented = false;

    private transient boolean pkBoundaryFound = false;


    public BigDatabaseExtractor(DatabaseExtractSetting setting) {
        super(setting);
        this.setName(setting.getName());

    }

    @Override
    public synchronized void init(JobContext context) {
        if (getSetting().isParallel()) {
            init_using_pkRange(context);
        } else
            init0(context);
    }


    public synchronized void init0(JobContext context) {
        getSetting().setDql(populate(getSetting().getDql(), context));
        if (!isReady()) {
            try {
                if (getSetting().getDatasource() != null) {
                    setDataSource(context.getResource(getSetting().getDatasource()));
                    setConnection(getDataSource().getConnection());
                    //get min and max of pk
                    this.pkStatement = getConnection().prepareStatement(getSetting().getPkDql());
                    this.pkResultSet = this.pkStatement.executeQuery();
                    if (this.pkResultSet.next()) {
                        this.minPk = this.pkResultSet.getObject(1);
                        this.maxPk = this.pkResultSet.getObject(2);
                        if (!getSetting().isUnion())
                            this.pkBoundary = this.minPk;
                    } else {
                        throw new EtltRuntimeException("Extractor init error: no pk range found for " + getName());
                    }

                    setStatement(getConnection().prepareStatement(getSetting().getDql()));
                    getStatement().setObject(1, this.minPk);
                    setResultSet(getStatement().executeQuery());
                    ResultSetMetaData resultSetMetaData = getResultSet().getMetaData();
                    int columnCount = resultSetMetaData.getColumnCount();
                    if (getSetting().isAutoResolve() && getColumns().size() == 0) {
                        for (int i = 0; i < columnCount; i++) {
                            getColumns().add(resultSetMetaData.getColumnLabel(i + 1));
                        }
                    }
                }
                setReady(true);
            } catch (SQLException e) {
                throw new EtltRuntimeException("Extractor init error: " + getName(), e);
            }
        }
    }

    public synchronized void init_using_pkRange(JobContext context) {
        getSetting().setDql(populate(getSetting().getDql(), context));
        if (!isReady()) {
            try {
                if (getSetting().getDatasource() != null) {
                    setDataSource(context.getResource(getSetting().getDatasource()));
                    setConnection(getDataSource().getConnection());
                    //get min and max of pk
                    this.pkStatement = getConnection().prepareStatement(getSetting().getPkDql());
                    this.pkResultSet = this.pkStatement.executeQuery();
                    if (this.pkResultSet.next()) {
                        this.minPk = this.pkResultSet.getObject(1);
                        this.maxPk = this.pkResultSet.getObject(2);
                        PkRange pkRange = new PkRange(Long.parseLong(String.valueOf(this.minPk)),
                                Long.parseLong(String.valueOf(this.maxPk)));
                        setPkFullRange(pkRange);
                        if (!getSetting().isUnion())
                            this.pkBoundary = this.minPk;
                    } else {
                        throw new EtltRuntimeException("Extractor init error: no pk range found for " + getName());
                    }

                    setStatement(getConnection().prepareStatement(getSetting().getDql()));
                    /*
                    use the same pk range to get result set meta data
                     */
                    getStatement().setObject(1, this.getPkFullRange().getStartPk());
                    getStatement().setObject(2, this.getPkFullRange().getStartPk());
                    setResultSet(getStatement().executeQuery());
                    ResultSetMetaData resultSetMetaData = getResultSet().getMetaData();
                    int columnCount = resultSetMetaData.getColumnCount();
                    if (getSetting().isAutoResolve() && getColumns().size() == 0) {
                        for (int i = 0; i < columnCount; i++) {
                            getColumns().add(resultSetMetaData.getColumnLabel(i + 1));
                        }
                    }
                }
                setReady(true);
            } catch (SQLException e) {
                throw new EtltRuntimeException("Extractor init error: " + getName(), e);
            }
        }
    }

    @Override
    public boolean isClosed() {
        return false;
    }

    @Override
    public void extract(JobContext context) {
        if (getSetting().isUnion()) {
            extract_unionPK(context);
        } else {
            extract_singlePK(context);
        }
    }


    /**
     * support union pk (under test)
     *
     * @param context
     */
    private void extract_unionPK(JobContext context) {
        if (!isReady())
            init(context);
        try {
            if (getResultSet() == null) {
                setConnection(getDataSource().getConnection());
                setStatement(getConnection().prepareStatement(getSetting().getDql()));
                getStatement().setObject(1, this.pkBoundary);
                // reset this pk boundary
                resetPkBoundary();
                setResultSet(getStatement().executeQuery());
            }
            if (getResultSet().next()) {
                if (this.skip < getSetting().getSkip()) {
                    this.skip++;
                    this.index++;
                    extract_unionPK(context);
                } else {
                    Map<String, Object> rowData = new HashMap<>();
                    ResultSetMetaData resultSetMetaData = getResultSet().getMetaData();
                    for (int i = 0; i < getColumns().size(); i++) {
                        if (getColumns().contains(resultSetMetaData.getColumnLabel(i + 1)))
                            rowData.put(getColumns().get(i),
                                    DatabaseUtil.getObject(getResultSet(), i + 1,
                                            resultSetMetaData.getColumnType(i + 1))
                            );
                    }
//                    Object pkValue =
//                            this.expressionCompiler.evaluate(getSetting().getPkName(), (name) ->
//                                    Variable.createVariable(name, rowData.get(name))
//                            );
                    Object pkValue = rowData.get(getSetting().getPkName());
                    if (pkValue == null)
                        throw new EtltRuntimeException("no pk value found: " + getSetting().getPkName());
                    //联合主键，使用联合主键时，查询时必须倒序排列
                    if (this.pkBoundary == null) {
                        this.pkBoundary = pkValue;
                        extract_unionPK(context);
                    } else {
                        if (this.pkBoundary.equals(pkValue)) {
                            extract_unionPK(context);
                        } else if (!pkBoundaryFound) {
                            this.pkBoundary = pkValue;
                            pkBoundaryFound = true;
                        }
                        Entity entity = new Entity(index++, rowData);
                        context.setEntity(this.getSetting().getName(), entity);
                    }
                }
            } else if (!this.pkBoundary.equals(this.maxPk) && !complemented) {
                getStatement().setObject(1, this.pkBoundary);
                setResultSet(getStatement().executeQuery());
                resetPkBoundary();
                extract_unionPK(context);
            } else if (getSetting().getDqlByPk() != null && !complemented) {
                //-- add the minimum data with minimum pk
                setStatement(getConnection().prepareStatement(getSetting().getDqlByPk()));
                getStatement().setObject(1, this.minPk);
                setResultSet(getStatement().executeQuery());
                extract_unionPK(context);
                complemented = true;
            } else {
                context.removeEntity(this.getSetting().getName());
            }
        } catch (Throwable e) {
            throw new EtltRuntimeException("etl exception caught where extract " + getName(), e);
        }
    }

    /**
     * support single pk (fully tested)
     *
     * @param context
     */
    private void extract_singlePK(JobContext context) {
        if (!(isReady()))
            init(context);
        try {
            if (getResultSet() == null) {
                setConnection(getDataSource().getConnection());
                setStatement(getConnection().prepareStatement(getSetting().getDql()));
                getStatement().setObject(1, this.pkBoundary);
                setResultSet(getStatement().executeQuery());
            }
            if (getResultSet().next()) {
                if (this.skip < getSetting().getSkip()) {
                    this.skip += 1;
                    this.index += 1;
                    extract(context);
                } else {
                    Map rowData = new HashMap();
                    ResultSetMetaData resultSetMetaData = getResultSet().getMetaData();
                    for (int i = 0; i < getColumns().size(); ++i) {
                        if (getColumns().contains(resultSetMetaData.getColumnLabel(i + 1))) {
                            rowData.put(getColumns().get(i),
                                    DatabaseUtil.getObject(getResultSet(), i + 1, resultSetMetaData
                                            .getColumnType(i + 1)));
                        }
                    }

                    Object pkValue = rowData.get(getSetting().getPkName());
                    if (pkValue == null)
                        throw new EtltRuntimeException(new StringBuilder().append("no pk value found: ").append(getSetting().getPkName()).toString());
                    this.pkBoundary = pkValue;
                    Entity entity = new Entity(this.index++, rowData);
                    context.setEntity(getSetting().getName(), entity);
                }
            } else if(this.pkBoundary == null){
                this.complemented = true;
            }else if (!this.pkBoundary.equals(this.maxPk) && !this.complemented) {
                getStatement().setObject(1, this.pkBoundary);
                setResultSet(getStatement().executeQuery());
                extract_singlePK(context);
            } else if (getSetting().getDqlByPk() != null && !this.complemented) {
                setStatement(getConnection().prepareStatement(getSetting().getDqlByPk()));
                getStatement().setObject(1, this.minPk); //todo
                setResultSet(getStatement().executeQuery());
                extract_singlePK(context);
                this.complemented = true;
            } else {
                context.removeEntity(getSetting().getName());
            }
        } catch (Throwable e) {
            throw new EtltRuntimeException(new StringBuilder().append("etl exception caught where extract ").append(getName()).toString(), e);
        }
    }


    @Override
    public BigDatabaseExtractor createInstance() {
        BigDatabaseExtractor extractor = new BigDatabaseExtractor(this.getSetting());
        extractor.setDataSource(getDataSource());
        return extractor;
    }

    private void resetPkBoundary() {
        this.pkBoundary = null;
        this.pkBoundaryFound = false;
    }

    @Override
    public PkRange getPkFullRange() {
        return this.pkFullRange;
    }
}
